import pandas as pd
import numpy as np
import torch

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 处理文件夹数据（重要）
def read_label(label_path='datasets/test/dmer_annotations(std).csv'):

    # 读取 CSV 文件
    labels = pd.read_csv(label_path)
    # 按照 musicId 分组，提取每首歌的 VA 值
    music_groups = labels.groupby('musicId')[['Arousal(std)', 'Valence(std)']]
    # 存储每首歌id，以便于读取音频文件
    music_id = np.array([music_id for music_id, _ in music_groups])
    # 将每首歌的 VA 值组合成一个二维数组
    va_matrix = np.array([group[['Arousal(std)', 'Valence(std)']].values for _, group in music_groups], dtype=object)
    return va_matrix, music_id  # va_matrix的shape为[nums_music, nums_va, [V, A]]


# 音频分帧
def frame_audio(waveform, frame_length=22050, frame_step=22050):
    frames = waveform.unfold(
        dimension=1,
        size=frame_length,
        step=frame_step
    )  # shape: [1, nums_frames, frame_length]
    # # 获取最后一帧的长度
    # last_frame_start = frames.shape[1] * frame_step
    # last_frame = waveform[:, last_frame_start:]
    # # 处理最后一帧
    # if last_frame.size(1) > frame_length / 2:
    #     # 如果最后一帧大于等于 frame_length / 2，则补齐
    #     padding = frame_length - last_frame.size(1)
    #     last_frame = torch.cat([last_frame, torch.zeros(1, padding)], dim=1)
    #     frames = torch.cat([frames, last_frame.unsqueeze(1)], dim=1)
    return frames


# 加窗函数
def windows(frames, frame_length):
    # 生成窗并确保在正确的设备上
    window = torch.hann_window(frame_length).to(frames.device)  # 使用与输入相同的设备
    window = window.unsqueeze(0).unsqueeze(0)  # 扩展维度
    frames_windowed = frames * window
    return frames_windowed


if __name__ == '__main__':
    # print(va_matrix)
    # 查看处理后的二维数组

    print(read_label()[0].shape)     # 所有歌的 VA 值数据
    print(read_label()[0][0].shape)  # 第一首歌的 VA 值数据
    print(type(read_label()[0]))
